import os

import streamlit as st
from pandasai import SmartDataframe
from pandasai.callbacks import BaseCallback
from pandasai.llm import OpenAI
from pandasai.responses.response_parser import ResponseParser
import pandas as pd
import numpy as np
from data import load_data
import os
# 禁用 SSL 验证（仅用于调试）
os.environ["REQUESTS_CA_BUNDLE"] = ""
os.environ["SSL_CERT_FILE"] = ""

class StreamlitCallback(BaseCallback):
    def __init__(self, container) -> None:
        """Initialize callback handler."""
        self.container = container

    def on_code(self, response: str):
        self.container.code(response)


class StreamlitResponse(ResponseParser):
    def __init__(self, context) -> None:
        super().__init__(context)

    def format_dataframe(self, result):
        st.dataframe(result["value"])
        return

    def format_plot(self, result):
        st.image(result["value"])
        return

    def format_other(self, result):
        st.write(result["value"])
        return


st.write("# Chat with Credit Card Fraud Dataset 🦙")

# 员工信息
data = {
    "员工ID": np.arange(1, 11),  # 员工编号
    "姓名": ["张三", "李四", "王五", "赵六", "孙七", "周八", "吴九", "郑十", "钱伯", "孔仲"],
    "性别": np.random.choice(["男", "女"], size=10),  # 随机分配性别
    "年龄": np.random.randint(20, 60, size=10),  # 随机生成年龄
    "职位": np.random.choice(["护林员", "导游", "管理员", "清洁工"], size=10),  # 随机分配职位
    "入职日期": pd.date_range(start="2015-01-01", periods=10, freq="M"),  # 随机生成入职日期
    "工资": np.random.randint(3000, 8000, size=10)  # 随机生成工资
}

# 创建DataFrame
df = pd.DataFrame(data)

with st.expander("🔎 Dataframe Preview"):
    st.write(df.tail(3))

query = st.text_area("🗣️ Chat with Dataframe")
container = st.container()

if query:
    llm = OpenAI(api_token="sk-lufpkjmnhkzdfuuosshszezhnjldzyvwvjseflxbatvtzmdl", openai_proxy="https://api.siliconflow.cn/v1")  

    query_engine = SmartDataframe(
        df,
        config={
            "llm": llm,
            "response_parser": StreamlitResponse,
            "callback": StreamlitCallback(container),
        },
    )

    answer = query_engine.chat(query)